package dan.backend
import chisel3._
import chisel3.util._ 
import dan.common._ 
import scala.collection.mutable.ArrayBuffer
import dan.frontend.BrUpdateInfo
import dan.utility.killedByMispredict
import dan.utility.UpdatedNewUop
import dan.common.Consts.REGT_GPR
import dan.utility.UpdatedBrMask

class RegisterFileReadPortIO (
    val addrWidth: Int, val dataWidth: Int
) extends CoreBundle {
    val addr = Input(UInt(addrWidth.W))
    val data = Output(UInt(dataWidth.W))
}

class RegisterFileWritePort (
    val addrWidth: Int, val dataWidth: Int
) extends CoreBundle {
    val addr = UInt(addrWidth.W)
    val data = UInt(dataWidth.W)
}

object WritePort {
    def apply(
        enq: DecoupledIO[ExecUnitResp], addrWidth: Int, dataWidth: Int, rtype: UInt
    ): Valid[RegisterFileWritePort] = {
        val wport = Wire(Valid(new RegisterFileWritePort(addrWidth, dataWidth)))

        enq.ready       := true.B
        wport.valid     := enq.valid && enq.bits.uop.archDstType === rtype
        wport.bits.addr := enq.bits.uop.physDst
        wport.bits.data := enq.bits.data
        wport
    }
}

abstract class RegisterFile (
    numRegisters: Int,
    numReadPorts: Int,
    numWritePorts: Int,
    registerWidth: Int,
    bypassableArray: Seq[Boolean]
) extends CoreModule {
    val io = IO(new CoreBundle{
        val read_ports  = Vec(
            numReadPorts, new RegisterFileReadPortIO(physRegBits, registerWidth)
        )
        val write_ports = Flipped(Vec(
            numWritePorts, Valid(new RegisterFileWritePort(physRegBits, registerWidth))
        ))
    })
}

class RegisterFileSynthesizable (
    numRegisters : Int,
    numReadPorts : Int,
    numWritePorts: Int,
    registerWidth: Int,
    bypassableArray: Seq[Boolean]
) extends RegisterFile(
    numRegisters,
    numReadPorts,
    numWritePorts,
    registerWidth,
    bypassableArray
) {
    val regfile = Mem(numRegisters,UInt(registerWidth.W))

    val read_addrs = io.read_ports.map(p => RegNext(p.addr)) // delay 1 cycle
    val read_data  = Wire(Vec(numReadPorts,UInt(registerWidth.W)))

    for (i <- 0 until numReadPorts) {
        read_data(i) := regfile(read_addrs(i))
    }

    if (bypassableArray.reduce(_||_)) {
        val bypassableWports = ArrayBuffer[Valid[RegisterFileWritePort]]()
        io.write_ports zip bypassableArray map {
            case (wport,b) => if (b) { bypassableWports += wport }
        }

        for ( i <- 0 until numReadPorts){
            val bypassEns = bypassableWports.map(
                x => x.valid && x.bits.addr === read_addrs(i)
            )
            val bypassData = Mux1H(
                VecInit(bypassEns.toSeq),
                VecInit(bypassableWports.map(_.bits.data).toSeq)
            )

            io.read_ports(i).data := Mux(bypassEns.reduce(_|_), bypassData, read_data(i))
        }
    } else {
        for (i <- 0 until numReadPorts) {
            io.read_ports(i).data := read_data(i)
        }
    }

    for (wport <- io.write_ports) {
        when (wport.valid) {
            regfile(wport.bits.addr) := wport.bits.data
        }
    }
}

//////////////////////////////

class RegisterRead (
    issueWidth: Int,
    supportedUnitsArray: Seq[OpFuncType],
    numReadPortsArray: Seq[Int],
    numTotalBypassPorts: Int,
    registerWidth: Int
) extends CoreModule {
    val io = IO(new Bundle{
        val iss_valids = Input(Vec(issueWidth, Bool()))
        val iss_uops   = Input(Vec(issueWidth, new UOp))

        val rf_read_ports  = Flipped(Vec(numReadPortsArray.sum, new RegisterFileReadPortIO(physRegBits, registerWidth)))

        val bypass   = Input(Vec(numTotalBypassPorts, Valid(new ExecUnitResp(registerWidth))))

        val exe_reqs = Vec(issueWidth, Decoupled(new FuncUnitReq))

        val kill     = Input(Bool())
        val brupdate = Input(new BrUpdateInfo())
    })
    io.exe_reqs.foreach(_.bits.kill := io.kill)

    var idx = 0
    for (w <- 0 until issueWidth) {
        val numReadPorts = numReadPortsArray(w)

        val rrd_valid = Reg(Bool())
        val rrd_uop   = Reg(new UOp)
        // clock 0: func decode
        val dec_unit = Module(new RegisterReadDecode(supportedUnitsArray(w)))
        dec_unit.io.iss_valid := io.iss_valids(w)
        dec_unit.io.iss_uop   := io.iss_uops(w)

        rrd_valid := (
            dec_unit.io.rrd_valid && !killedByMispredict(io.brupdate,dec_unit.io.rrd_uop)
        )
        rrd_uop   := UpdatedNewUop(dec_unit.io.rrd_uop, io.brupdate)

        val rrd_rs1  = Wire(UInt(registerWidth.W))
        val rrd_rs2  = if (numReadPorts == 2) Wire(UInt(registerWidth.W)) else null
        val rs1_addr = if (numReadPorts == 2) {
            io.iss_uops(w).physOprand1
        } else {
            Mux(io.iss_uops(w).archOp2Type === REGT_GPR, io.iss_uops(w).physOprand2, io.iss_uops(w).physOprand1)
        }
        val rs2_addr = if (numReadPorts == 2) io.iss_uops(w).physOprand2 else null
        // clock 0: send read request to register file
        io.rf_read_ports(idx + 0).addr := rs1_addr
        if (numReadPorts == 2) io.rf_read_ports(idx + 1).addr := rs2_addr
        // clock 1: read data from register file
        rrd_rs1 := Mux(
            RegNext(rs1_addr === 0.U), 0.U, io.rf_read_ports(idx + 0).data
        )
        if (numReadPorts == 2) rrd_rs2 := Mux(
            RegNext(rs2_addr === 0.U), 0.U, io.rf_read_ports(idx + 1).data
        )
        idx += numReadPorts

        // clock 1: do bypass
        val bypassed_rs1 = Wire(UInt(registerWidth.W))
        val bypassed_rs2 = if (numReadPorts == 2) Wire(UInt(registerWidth.W)) else null
        val prs1 = if (numReadPorts == 2) {
            rrd_uop.physOprand1
        } else {
            Mux(rrd_uop.archOp2Type === REGT_GPR, rrd_uop.physOprand2, rrd_uop.physOprand1)
        }
        val lrs1Rtype = if (numReadPorts == 2) {
            rrd_uop.archOp1Type
        } else {
            Mux(rrd_uop.archOp2Type === REGT_GPR, rrd_uop.archOp2Type, rrd_uop.archOp1Type)
        }
        val prs2 = if (numReadPorts == 2) rrd_uop.physOprand2 else null
        val lrs2Rtype = if (numReadPorts == 2) rrd_uop.archOp2Type else null
        var rs1_cases = Seq((false.B, 0.U(registerWidth.W)))
        var rs2_cases = if (numReadPorts == 2) Seq((false.B, 0.U(registerWidth.W))) else null
        io.bypass foreach { bypass =>
            rs1_cases ++= Seq((
                bypass.valid && (prs1 === bypass.bits.uop.physDst) && bypass.bits.uop.regWrEnable &&
                bypass.bits.uop.archDstType === REGT_GPR && lrs1Rtype === REGT_GPR && (prs1 =/= 0.U),
                bypass.bits.data
            ))
            if (numReadPorts == 2) {
                rs2_cases ++= Seq((
                    bypass.valid && (prs2 === bypass.bits.uop.physDst) && bypass.bits.uop.regWrEnable &&
                    bypass.bits.uop.archDstType === REGT_GPR && lrs2Rtype === REGT_GPR && (prs2 =/= 0.U),
                    bypass.bits.data
                ))
            }
        }
        bypassed_rs1 := MuxCase(rrd_rs1, rs1_cases)
        if (numReadPorts == 2) bypassed_rs2 := MuxCase(rrd_rs2, rs2_cases)

        // clock 2: send request to execution units
        val exe_req_valid = RegInit(false.B)
        val exe_req_uop   = Reg(new UOp)
        val exe_req_rs1   = Reg(UInt(registerWidth.W))
        val exe_req_rs2   = if (numReadPorts == 2) Reg(UInt(registerWidth.W)) else null

        val killed = io.kill || killedByMispredict(io.brupdate, rrd_uop)
        exe_req_valid      := Mux(killed, false.B, rrd_valid)
        exe_req_uop        := rrd_uop
        exe_req_uop.brMask := UpdatedBrMask(rrd_uop, io.brupdate)

        exe_req_rs1 := bypassed_rs1
        if (numReadPorts == 2) exe_req_rs2 := bypassed_rs2

        io.exe_reqs(w).valid        := exe_req_valid
        io.exe_reqs(w).bits.uop     := exe_req_uop
        io.exe_reqs(w).bits.oprand1 := exe_req_rs1
        io.exe_reqs(w).bits.oprand2 := (if (numReadPorts == 2) exe_req_rs2 else DontCare)
    }
}



